﻿#pragma once

#include <vector>
#include <wbemcli.h>
#include <atlbase.h>

#include "com_init.h"
#include "com_interface.h"
#include "process_query_info.h"

class process_query
{
public:
    /**
     * @brief 查询前的准备工作
     */
    void prepare();

    /**
     * @brief 发起WMI查询请求
     *
     * @return 查询到的进程信息
     */
    std::vector<process_query_info> execute_query();

private:
    /**
     * @brief 设置COM安全信息
     */
    void set_security();

    /**
     * @brief 创建进程内COM服务器
     */
    void create_instance();

    /**
     * @brief 连接WMI命名空间
     */
    void connect_server();

    /**
     * @brief 设置WMI连接的安全等级
     */
    void set_proxy_blanket();

    /**
     * @brief 提取查询结果
     *
     * @param enumerator 枚举对象
     * @return 查询结果
     */
    std::vector<process_query_info> fetch_results(
        com_interface<IEnumWbemClassObject> enumerator);

    /**
     * @brief 获取下一个进程信息
     *
     * @param enumerator 枚举对象
     * @param vec 存放查询结果信息的容器
     * @return 获取成功则返回TRUE，否则返回FALSE
     */
    BOOL get_next_process_info(
        IEnumWbemClassObject* enumerator,
        std::vector<process_query_info>& vec);

private:
    com_init init;  // COM初始化和反初始化
    com_interface<IWbemLocator> locator;
    com_interface<IWbemServices> service;
};
